from __future__ import annotations

import os, pathlib, itertools
from typing import Mapping
from google.protobuf.compiler import plugin_pb2 as plugin
from google.protobuf.descriptor_pb2 import (
    DescriptorProto as MessageDP,
    EnumDescriptorProto as EnumDP,
    FieldDescriptorProto as FieldDP,
    FileDescriptorProto as FileDP,
    ServiceDescriptorProto as ServiceDP,
    OneofDescriptorProto as OneofDP,
)
from google.protobuf.descriptor_pool import DescriptorPool
from .options import parse_options
from .templates import TemplateRenderer
from .descriptor_ext import ProtoFile
from .utils import split_package_and_type_name_of_field_type

# global options
g_enable_namespace = False


def generate_single_file_dts(proto: ProtoFile, file_map: Mapping[str, ProtoFile], renderer: TemplateRenderer) -> str:
    global g_enable_namespace
    g_enable_namespace = proto.opts.enable_namespace
    namespace_levels: list[str] = []
    if g_enable_namespace:
        namespace_levels = proto.fd.package.split(".")
    lines: list[str] = []

    # 1. 文件头
    lines.append("// Generated by protoc-gen-napi. DO NOT EDIT!")
    lines.append("// source: " + proto.fd.name)
    lines.append("")

    # 2. 依赖 import
    if not g_enable_namespace:
        type_dependencies = _collect_type_dependencies(proto.fd, renderer.pool)
        dependencies = {}
        for msg_dep in itertools.chain(type_dependencies["messages"], type_dependencies["enums"]):
            filename, _, type_name = msg_dep.split("|")
            dependencies.setdefault(filename, []).append(type_name.replace(".", "_"))
        for file_name, type_names in dependencies.items():
            rel_path = os.path.relpath(file_map[file_name].cc_dts_name[:-5], pathlib.Path(proto.fd.name).parent.as_posix())
            rel_path = rel_path if rel_path.startswith("../") else f"./{rel_path}"
            lines.append(f'import {{ {", ".join(type_names)} }} from "{rel_path}";')
        lines.append("")

    if g_enable_namespace:
        for level in namespace_levels:
            lines.append(f"namespace {level} {{")
        lines.append("")

    # 3. 枚举
    for enum in proto.fd.enum_type:
        lines.extend(_generate_enum_dts(enum, renderer, pkg_name=proto.fd.package))
        lines.append("")

    # 4. 消息
    for msg in proto.fd.message_type:
        lines.extend(
            _generate_message_dts(
                msg,
                renderer,
                package_name=proto.fd.package,
            )
        )
        lines.append("")

    if g_enable_namespace:
        for level in reversed(namespace_levels):
            lines.append(f"}} // namespace {level}")
        lines.append("")
        # lines.extend(_generate_single_file_typealias(proto))

    # 5. 服务
    for svc in proto.fd.service:
        lines.extend(_generate_service_dts(svc, renderer))
        lines.append("")

    # 6. 文件级扩展
    for ext, _ in proto.iter_extensions():
        lines.extend(_generate_extension_dts(ext, renderer))
        lines.append("")

    # 7. 文件级选项
    if proto.fd.options.ByteSize():
        lines.extend(_generate_file_options_dts(proto.fd.options, renderer))

    return "\n".join(lines)


# ---------------- 枚举 ----------------
def _generate_enum_dts(
    enum: EnumDP,
    renderer: TemplateRenderer,
    parent_name: str = "",
    pkg_name: str = "",
) -> list[str]:
    return renderer.enum_stub_dts(enum, parent_name)


# ---------------- 消息 ----------------
def _generate_message_dts(
    msg: MessageDP,
    renderer: TemplateRenderer,
    parent_names: list[str] = [],
    package_name: str = "",
) -> list[str]:
    message_name = "_".join([*parent_names, msg.name]) if parent_names and not g_enable_namespace else msg.name
    lines = []
    nested_messages = [nt for nt in msg.nested_type if not nt.options.map_entry]
    need_msg_namespace: bool = g_enable_namespace and len(msg.enum_type) + len(nested_messages) > 0
    if need_msg_namespace:
        lines.append(f'namespace {msg.name} {{')
    # 嵌套枚举
    for nested_enum in msg.enum_type:
        lines.extend(
            _generate_enum_dts(
                nested_enum,
                renderer,
                parent_name="" if g_enable_namespace else message_name,
                pkg_name=package_name,
            )
        )

    # 嵌套消息
    for nested_msg in nested_messages:
        lines.extend(_generate_message_dts(nested_msg, renderer, parent_names=[*parent_names, msg.name], package_name=package_name))
        lines.append("")

    if need_msg_namespace:
        lines.append(f"}} // namespace {msg.name}")
        lines.append("")

    parent_name = "_".join(parent_names) if parent_names and not g_enable_namespace else ''
    lines.extend(renderer.message_stub_dts(msg, parent_name, package_name))

    return lines


# ---------------- 服务 ----------------
def _generate_service_dts(svc: ServiceDP, renderer: TemplateRenderer) -> list[str]:
    return renderer.service_stub(svc)
    # TODO(user): 填充 service_stub 模板


# ---------------- 扩展 ----------------
def _generate_extension_dts(ext: FieldDP, renderer: TemplateRenderer) -> list[str]:
    return renderer.extension_stub(ext)
    # TODO(user): 填充 extension_stub 模板


# ---------------- 文件选项 ----------------
def _generate_file_options_dts(options, renderer: TemplateRenderer) -> list[str]:
    return renderer.file_options_stub(options)
    # TODO(user): 填充 file_options_stub 模板


def all_symbols_in_file(proto: ProtoFile) -> list[str]:
    symbols = [enum.name for enum in proto.fd.enum_type]

    def collect_msg_symbols(msg: MessageDP, parent_name: str = ""):
        if msg.options.map_entry:
            return
        msg_name = f"{parent_name}_{msg.name}" if parent_name else msg.name
        symbols.append(msg_name)
        symbols.extend([f"{msg_name}_{enum.name}" for enum in msg.enum_type])
        for nested_type in msg.nested_type:
            collect_msg_symbols(nested_type, msg_name)

    for msg in proto.fd.message_type:
        collect_msg_symbols(msg)
    return symbols


def all_symbols_in_file_ns(proto: ProtoFile) -> list[str]:
    if not g_enable_namespace:
        return []
    pkg_prefix = proto.fd.package.split(".")
    symbols = ["_".join([*pkg_prefix, enum.name]) for enum in proto.fd.enum_type]
    symbols.extend(["_".join([*pkg_prefix, "$" + msg.name]) for msg in proto.fd.message_type if not msg.options.map_entry])

    def collect_nested_symbols(msg: MessageDP, parent_names: list[str] = []):
        if msg.options.map_entry:
            return
        prefix = [*proto.fd.package.split("."), *parent_names, msg.name]
        symbols.extend(["_".join([*prefix, enum.name]) for enum in msg.enum_type])
        for nested_type in msg.nested_type:
            if nested_type.options.map_entry:
                continue
            symbols.append("_".join([*prefix, nested_type.name]))
            collect_nested_symbols(nested_type, [*parent_names, msg.name])

    for msg in proto.fd.message_type:
        collect_nested_symbols(msg)
    return symbols


def _generate_single_file_typealias(proto: ProtoFile) -> list[str]:
    if not g_enable_namespace:
        return []
    pkg_prefix = proto.fd.package.split(".")
    lines = []
    if proto.fd.package:
        for enum in proto.fd.enum_type:
            symbol_parts = [*pkg_prefix, enum.name]
            lines.append(f"export type {'_'.join(symbol_parts)} = {'.'.join(symbol_parts)}")

        for msg in proto.fd.message_type:
            if msg.options.map_entry:
                continue
            symbol_parts = [*pkg_prefix, msg.name]
            lines.append(f"export type {'_'.join(symbol_parts)} = {'.'.join(symbol_parts)}")
            symbol_parts = [*pkg_prefix, "$" + msg.name]
            lines.append(f"export type {'_'.join(symbol_parts)} = {'.'.join(symbol_parts)}")

    def collect_nested_symbols(msg: MessageDP, parent_names: list[str] = []):
        if msg.options.map_entry:
            return
        prefix = [*proto.fd.package.split("."), *parent_names, msg.name]
        for enum in msg.enum_type:
            lines.append(f"export type {'_'.join([*prefix, enum.name])} = {'.'.join([*prefix, enum.name])}")
        for nested_type in msg.nested_type:
            if nested_type.options.map_entry:
                continue
            symbol_parts = [*prefix, nested_type.name]
            lines.append(f"export type {'_'.join(symbol_parts)} = {'.'.join(symbol_parts)}")
            symbol_parts = [*prefix, "$" + nested_type.name]
            lines.append(f"export type {'_'.join(symbol_parts)} = {'.'.join(symbol_parts)}")
            collect_nested_symbols(nested_type, [*parent_names, msg.name])

    for msg in proto.fd.message_type:
        collect_nested_symbols(msg)
    return lines


def _collect_type_dependencies(file_proto: FileDP, pool: DescriptorPool):
    """
    返回结构：
    {
        "messages": ["file_name.proto|pkg_name|MsgA", "file_name.proto|pkg_name|MsgA"],
        "enums": ["file_name.proto|pkg_name|EnumA"]
    }
    """
    result = {}

    def add_dep(file_name: str, kind: str, pkg_name: str, type_name: str):
        if file_name != file_proto.name:
            result.setdefault(kind, set()).add(f"{file_name}|{pkg_name}|{type_name}")

    # ---- 遍历 message ----
    def visit_message(msg_proto: MessageDP, prefix: str):
        """prefix = 包名 + "." + message名（用于输出全名）"""

        # 遍历字段
        for field in msg_proto.field:
            if field.type in (FieldDP.TYPE_MESSAGE, FieldDP.TYPE_ENUM):
                type_name = field.type_name.lstrip(".")  # 去掉前导点

                # 从 pool 中找真正的 Descriptor
                try:
                    if field.type == FieldDP.TYPE_MESSAGE:
                        desc = pool.FindMessageTypeByName(type_name)
                        pkg_name, type_name = split_package_and_type_name_of_field_type(field, pool)
                        add_dep(desc.file.name, "messages", pkg_name, type_name)
                    else:
                        desc = pool.FindEnumTypeByName(type_name)
                        pkg_name, type_name = split_package_and_type_name_of_field_type(field, pool)
                        add_dep(desc.file.name, "enums", pkg_name, type_name)
                except KeyError:
                    # 有可能 pool 中没有某些描述
                    pass

        # 递归嵌套 message
        for nested in msg_proto.nested_type:
            nested_full = f"{prefix}.{nested.name}"
            visit_message(nested, nested_full)

    # ---- 遍历所有顶层 message ----
    pkg = file_proto.package
    for msg in file_proto.message_type:
        name = f"{pkg}.{msg.name}" if pkg else msg.name
        visit_message(msg, name)

    # ---- 将 set 转为 list ----
    result["messages"] = sorted(result.get("messages", []))
    result["enums"] = sorted(result.get("enums", []))
    return result
